
use super::Var;

#[derive(Debug, PartialEq, Clone)]
pub enum Operate {
    Binary(BinaryOperate),
    Unary(UnaryOperate)
}

#[derive(Debug, PartialEq, Clone)]
pub enum BinaryOperate {
    Add(Var, Var),
    Mul(Var, Var),
    Sub(Var, Var),
    Div(Var, Var),
}

impl BinaryOperate {
    pub fn backward(&self, var: &Var) {
        match self {
            Self::Add(left, right) => {
                left.increase_grad(var.gard());
                right.increase_grad(var.gard());
            }
            Self::Mul(left, right) => {
                left.increase_grad(right.data() * var.gard());
                right.increase_grad(left.data() * var.gard());
            }
            Self::Sub(left, right) => {
                left.increase_grad(var.gard());
                right.increase_grad(-var.gard());
            }
            Self::Div(left, right) => {
                left.increase_grad(var.gard() * (1.0 / right.data()));
                right.increase_grad(var.gard() * (-left.data() * right.data().powf(-2.)));
            }
        }
    }

    pub fn operands(&self) -> (&Var, &Var) {
        match self {
            Self::Add(left, right) => (left, right),
            Self::Mul(left, right) => (left, right),
            Self::Sub(left, right) => (left, right),
            Self::Div(left, right) => (left, right),
        }
    }
}

#[derive(Debug, PartialEq, Clone)]
pub enum UnaryOperate {
    Pow(Var, f64),
    Tanh(Var),
    Exp(Var),
    Log(Var, f64),
    Abs(Var),
    Relu(Var),
}

impl UnaryOperate {
    pub fn backward(&self, var: &Var) {
        match self {
            Self::Pow(base, power) => {
                let grad = power * base.data().powf(power - 1.) * var.gard();
                base.increase_grad(grad);
            }
            Self::Tanh(prev) => {
                let grad = (1.0 - var.data().powf(2.0)) * var.gard();
                prev.increase_grad(grad);
            }
            Self::Exp(prev) => {
                let grad = var.data() * var.gard();
                prev.increase_grad(grad);
            }
            Self::Log(prev, base) => {
                // prev.log(base) = var
                let grad = (1. / (prev.data() * base.ln())) * var.gard();
                prev.increase_grad(grad);
            }
            Self::Abs(prev) => {
                let grad = match prev.data() {
                    0. => 0.,
                    v if v > 0. => 1.,
                    _ => -1.,
                } * var.gard();
                prev.increase_grad(grad);
            }
            Self::Relu(prev) => {
                let grad = if var.data() > 0. { var.gard() } else { 0. };
                prev.increase_grad(grad);
            }
        }
    }

    pub fn operand(&self) -> &Var {
        match self {
            Self::Pow(base, _) => base,
            Self::Tanh(prev) => prev,
            Self::Exp(prev) => prev,
            Self::Log(prev, _) => prev,
            Self::Abs(prev) => prev,
            Self::Relu(prev) => prev,
        }
    }
}
